import numpy as np
import torch
import cv2
import random
import json
import folder_paths

from PIL import Image, ImageDraw
from nodes import SaveImage
from .utils import *



class PointPreview(SaveImage):
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "points": ("Point",),
                "image_width": ("INT", {"default": 512, "min": 1, "max": 10000, "step": 1}),
                "image_height": ("INT", {"default": 512, "min": 1, "max": 10000, "step": 1}),
            }
        }

    FUNCTION = "generate_point_preview"
    CATEGORY = "🌟FJDH/Point tool"
    OUTPUT_NODE = True

    def __init__(self):
        self.output_dir = folder_paths.get_temp_directory()
        self.type = "temp"
        self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz1234567890") for x in range(5))
        self.compress_level = 4

    def generate_point_preview(self, points, image_width, image_height):
        radius = min(image_width, image_height) // 250
        if radius < 1:
            radius = 1

        color = "white"

        masks = []

        for batch_points in points:
            mask = Image.new("RGB", (image_width, image_height), "black")
            draw = ImageDraw.Draw(mask)

            for point in batch_points:
                x, y = point
                draw.ellipse([x - radius, y - radius, x + radius, y + radius], fill=color)

            masks.append(mask)

        mask_tensor = pil2tensor(masks)
        mask = mask_tensor[:, :, :, 0]

        if mask.dim() == 2:
            mask = torch.unsqueeze(mask, 0)
        
        preview = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)

        return self.save_images(preview, "PointPreview")



###########################################################################################

class MaskToPoint:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "mask": ("MASK",),
                "num_centroids": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
                "top_points_ratio": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01}),
                "num_points": ("INT", {"default": 1000, "min": 1, "max": 1000, "step": 1}),
                "calculate_centroid": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ("Point", "BBOX")
    RETURN_NAMES = ("Points", "Bboxes")
    FUNCTION = "process_mask"
    CATEGORY = "🌟FJDH/Point tool"

    def process_mask(self, mask, num_centroids, top_points_ratio, num_points, calculate_centroid):
        all_points = []
        all_bboxes = []

        for batch_mask in mask:
            mask_np = batch_mask.squeeze().cpu().numpy()
            mask_binary = (mask_np > 0).astype(np.uint8)

            contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            polygons = [cv2.approxPolyDP(cnt, 0.01 * cv2.arcLength(cnt, True), True) for cnt in contours]

            centroids = []
            bboxes = []
            points_with_bbox = []

            for polygon in polygons:
                M = cv2.moments(polygon)
                if M["m00"] != 0:
                    cX = int(M["m10"] / M["m00"])
                    cY = int(M["m01"] / M["m00"])
                    centroids.append((cX, cY))

                    rect = cv2.minAreaRect(polygon)
                    box = cv2.boxPoints(rect)
                    box = np.intp(box)

                    min_x = min(box[:, 0])
                    min_y = min(box[:, 1])
                    max_x = max(box[:, 0])
                    max_y = max(box[:, 1])
                    bbox = (min_x, min_y, max_x, max_y)
                    bboxes.append(bbox)

                    if num_centroids == 1:
                        points_with_bbox.append(((cX, cY), bbox))
                    else:
                        points = []
                        min_x = min(polygon[:, 0, 0])
                        max_x = max(polygon[:, 0, 0])
                        min_y = min(polygon[:, 0, 1])
                        max_y = max(polygon[:, 0, 1])
                        while len(points) < num_points:
                            x = np.random.randint(min_x, max_x + 1)
                            y = np.random.randint(min_y, max_y + 1)
                            if cv2.pointPolygonTest(polygon, (x, y), False) >= 0:
                                points.append((x, y))

                        distance_map = cv2.distanceTransform(mask_binary, cv2.DIST_L2, 5)

                        points_with_distance = [(point, distance_map[point[1], point[0]]) for point in points]
                        points_with_distance.sort(key=lambda x: x[1], reverse=True)

                        num_top_points = int(len(points_with_distance) * top_points_ratio)
                        top_points_with_distance = points_with_distance[:num_top_points]

                        interval_size = num_top_points // num_centroids
                        optimized_points = []
                        selected_points = []

                        for i in range(num_centroids):
                            start_index = i * interval_size
                            end_index = start_index + interval_size if i < num_centroids - 1 else num_top_points
                            interval_points = top_points_with_distance[start_index:end_index]

                            if selected_points:
                                selected_points_array = np.array(selected_points)
                                distances = np.linalg.norm(selected_points_array[:, np.newaxis] - np.array([point[0] for point in interval_points]), axis=2)
                                min_distances = np.min(distances, axis=0)
                                farthest_point_index = np.argmax(min_distances)
                                selected_point = interval_points[farthest_point_index][0]
                            else:
                                selected_point = interval_points[np.random.randint(len(interval_points))][0]

                            selected_points.append(selected_point)
                            optimized_points.append(selected_point)

                        for point in optimized_points:
                            points_with_bbox.append((point, bbox))

            unique_points_with_bbox = {json.dumps(point): (point, bbox) for point, bbox in points_with_bbox}.values()

            unique_points = [point for point, bbox in unique_points_with_bbox]
            unique_bboxes = [bbox for point, bbox in unique_points_with_bbox]

            if calculate_centroid:
                if unique_points:
                    all_x = [point[0] for point in unique_points]
                    all_y = [point[1] for point in unique_points]
                    center_x = sum(all_x) // len(all_x)
                    center_y = sum(all_y) // len(all_y)
                    center_coordinates = [(center_x, center_y)]
                else:
                    center_coordinates = []
            else:
                center_coordinates = unique_points

            all_points.append(center_coordinates)
            all_bboxes.append(unique_bboxes)

        return (all_points, all_bboxes)

class KeypointSelector:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "openpose_json": ("POSE_KEYPOINT",),
                "keypoint_type": (["pose_keypoints_2d", "face_keypoints_2d", "hand_left_keypoints_2d", "hand_right_keypoints_2d"],),
                "indices": ("STRING", {"default": "0,1,30"}),
            },
        }

    RETURN_TYPES = ("Point",)
    RETURN_NAMES = ("Points",)
    FUNCTION = "select_keypoints"
    CATEGORY = "🌟FJDH/Point tool"

    def select_keypoints(self, openpose_json, keypoint_type, indices):
        indices = list(map(int, indices.split(',')))

        all_points = []

        for batch in openpose_json:
            keypoints = batch['people'][0][keypoint_type]
            selected_keypoints = [keypoints[i*3:i*3+3] for i in indices]
            points = [(int(point[0]), int(point[1])) for point in selected_keypoints]
            all_points.append(points)

        return (all_points,)


class GridPointGenerator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "width": ("INT", {"default": 512, "min": 1, "max": 10000, "step": 1}),
                "height": ("INT", {"default": 512, "min": 1, "max": 10000, "step": 1}),
                "rows": ("INT", {"default": 10, "min": 1, "max": 1000, "step": 1}),
                "cols": ("INT", {"default": 10, "min": 1, "max": 1000, "step": 1}),
                "batch_size": ("INT", {"default": 1, "min": 1, "max": 1000, "step": 1}),
            },
        }

    RETURN_TYPES = ("Point",)
    RETURN_NAMES = ("GridPoints",)
    FUNCTION = "generate_grid_points"
    CATEGORY = "🌟FJDH/Point tool"

    def generate_grid_points(self, width, height, rows, cols, batch_size):
        all_points = []

        x_step = (width - 2 * (width // (cols + 1))) // (cols - 1)
        y_step = (height - 2 * (height // (rows + 1))) // (rows - 1)

        start_x = width // (cols + 1)
        start_y = height // (rows + 1)

        for _ in range(batch_size):
            points = []
            for row in range(rows):
                for col in range(cols):
                    x = start_x + col * x_step
                    y = start_y + row * y_step
                    points.append((x, y))
            all_points.append(points)

        return (all_points,)




###########################################################################################

class CenterPointCalculator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "points": ("Point",),
            },
        }

    RETURN_TYPES = ("Point",)
    RETURN_NAMES = ("CenterPoint",)
    FUNCTION = "calculate_center"
    CATEGORY = "🌟FJDH/Point tool"


    def calculate_center(self, points):
        all_centers = []

        for batch_points in points:
            if batch_points:
                all_x = [point[0] for point in batch_points]
                all_y = [point[1] for point in batch_points]
                center_x = sum(all_x) // len(all_x)
                center_y = sum(all_y) // len(all_y)
                center_point = [(center_x, center_y)]
            else:
                center_point = ()

            all_centers.append(center_point)

        return (all_centers,)

class PointMerger:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "points1": ("Point",),
                "points2": ("Point",),
            },
            "optional": {
                "points3": ("Point",),
            }
        }

    RETURN_TYPES = ("Point",)
    RETURN_NAMES = ("MergedPoints",)
    FUNCTION = "merge_coordinates"
    CATEGORY = "🌟FJDH/Point tool"

    def merge_coordinates(self, points1, points2, points3=None):
        all_merged_points = []

        for batch1, batch2 in zip(points1, points2):
            merged_points = batch1 + batch2

            if points3:
                batch3 = points3.pop(0)
                merged_points += batch3

            all_merged_points.append(merged_points)

        return (all_merged_points,)


class PointExtractor:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "points": ("Point",),
                "index": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
            },
        }

    RETURN_TYPES = ("INT", "INT")
    RETURN_NAMES = ("x_INT", "y_INT")
    FUNCTION = "extract_coordinates"
    CATEGORY = "🌟FJDH/Point tool"

    def extract_coordinates(self, points, index):
        all_x = []
        all_y = []

        for batch_points in points:
            if 0 <= index < len(batch_points):
                x, y = batch_points[index]
            else:
                x = 0
                y = 0
            all_x.append(x)
            all_y.append(y)

            # 调试输出
            print(f"Batch points: {batch_points}")
            print(f"Index: {index}")
            print(f"Extracted point: ({x}, {y})")

        # 调试输出
        print(f"All x coordinates: {all_x}")
        print(f"All y coordinates: {all_y}")

        return (all_x, all_y)

class PointMover:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "points": ("Point",),
                "distance": ("INT", {"default": 10, "min": 0, "max": 10000, "step": 1}),
                "direction": (["up", "down", "left", "right"],),
            },
            "optional": {
                "multiple_points": ("Point",),
            },
        }

    RETURN_TYPES = ("Point",)
    RETURN_NAMES = ("MovedPoints",)
    FUNCTION = "move_coordinates"
    CATEGORY = "🌟FJDH/Point tool"

    def move_coordinates(self, points, distance, direction, multiple_points=None):
        all_moved_points = []

        if multiple_points is None:
            multiple_points = [None] * len(points)

        for batch_points, batch_multiple_points in zip(points, multiple_points):
            move_vector = np.array([0, 0])

            if batch_multiple_points:
                if len(batch_multiple_points) < 2:
                    raise ValueError("Multiple points must contain at least two points.")
                start_point = np.array(batch_multiple_points[0])
                end_point = np.array(batch_multiple_points[-1])
                vector = end_point - start_point
                unit_vector = vector / np.linalg.norm(vector)
                move_vector = unit_vector * distance
            else:
                if direction == "up":
                    move_vector[1] -= distance
                elif direction == "down":
                    move_vector[1] += distance
                elif direction == "left":
                    move_vector[0] -= distance
                elif direction == "right":
                    move_vector[0] += distance

            moved_coordinates = []
            for point in batch_points:
                point_array = np.array(point)
                moved_point = point_array + move_vector
                moved_coordinates.append((int(moved_point[0]), int(moved_point[1])))

            all_moved_points.append(moved_coordinates)

        return (all_moved_points,)



class PointReversor:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "points": ("Point",),
            },
        }

    RETURN_TYPES = ("Point",)
    RETURN_NAMES = ("ReversedPoints",)
    FUNCTION = "reverse_coordinates"
    CATEGORY = "🌟FJDH/Point tool"

    def reverse_coordinates(self, points):
        all_reversed_points = []

        for batch_points in points:
            reversed_points = batch_points[::-1]
            all_reversed_points.append(reversed_points)

        return (all_reversed_points,)

class DistanceCalculator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "points": ("Point",),
            },
        }

    RETURN_TYPES = ("INT",)
    RETURN_NAMES = ("Distance",)
    FUNCTION = "calculate_distance"
    CATEGORY = "🌟FJDH/Point tool"

    def calculate_distance(self, points):
        all_distances = []

        for batch_points in points:
            if len(batch_points) < 2:
                raise ValueError("Points must contain at least two points.")

            start_point = np.array(batch_points[0])
            end_point = np.array(batch_points[-1])

            distance = int(np.linalg.norm(end_point - start_point))
            all_distances.append(distance)

        return (all_distances,)



class PointThresholdFilter:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "points": ("Point",),
                "threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
            },
        }

    RETURN_TYPES = ("Point",)
    RETURN_NAMES = ("FilteredPoints",)
    FUNCTION = "filter_points"
    CATEGORY = "🌟FJDH/Point tool"

    def filter_points(self, points, threshold):
        all_filtered_points = []

        for batch_points in points:
            num_points_to_keep = int(len(batch_points) * threshold)
            filtered_points = random.sample(batch_points, num_points_to_keep)
            all_filtered_points.append(filtered_points)

        return (all_filtered_points,)

###########################################################################################

class PointToCoordinates:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "points": ("Point",),
            },
        }

    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("Coordinates", "Coordinate")
    FUNCTION = "point_to_coordinates"
    CATEGORY = "🌟FJDH/Coordinate tool"

    def point_to_coordinates(self, points):
        all_coordinates = []
        single_coordinate = []

        for batch_points in points:
            coordinates = [{"x": point[0], "y": point[1]} for point in batch_points]
            all_coordinates.append(coordinates)

        for batch_points in points:
            coordinates = [{"x": point[0], "y": point[1]} for point in batch_points]
            single_coordinate.append(coordinates)

        all_coordinates_json = json.dumps(all_coordinates)
        single_coordinate_json = json.dumps(single_coordinate[0])

        return (all_coordinates_json, single_coordinate_json)



class CoordinatesToPoint:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "coordinates": ("STRING",),
            },
        }

    RETURN_TYPES = ("Point",)
    RETURN_NAMES = ("Points",)
    FUNCTION = "coordinates_to_point"
    CATEGORY = "🌟FJDH/Coordinate tool"

    def coordinates_to_point(self, coordinates):
        coordinates = json.loads(coordinates)
        
        # 如果输入是一维列表，将其转换为二维列表
        if not isinstance(coordinates[0], list):
            coordinates = [coordinates]
        
        # 如果输入是二维列表，将其转换为三维列表
        if not isinstance(coordinates[0][0], list):
            coordinates = [coordinates]
        
        points = [[(point['x'], point['y']) for point in batch] for batch in coordinates]
        return (points,)  # 返回一个包含单个批次的列表

    

class AngleCalculator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "points": ("Point",),
            },
        }

    RETURN_TYPES = ("FLOAT", "FLOAT")
    RETURN_NAMES = ("Angle", "DirectionAngle")
    FUNCTION = "calculate_angle"
    CATEGORY = "🌟FJDH/Point tool"

    def calculate_angle(self, points):
        all_angles = []
        all_direction_angles = []

        for batch_points in points:
            if len(batch_points) < 2:
                raise ValueError("Points must contain at least two points.")

            start_point = np.array(batch_points[0])
            end_point = np.array(batch_points[-1])

            vector = end_point - start_point
            horizontal_vector = np.array([1, 0])
            angle_rad = np.arccos(np.dot(vector, horizontal_vector) / (np.linalg.norm(vector) * np.linalg.norm(horizontal_vector)))
            angle_deg = np.degrees(angle_rad)

            if vector[1] < 0:
                angle_deg = 360 - angle_deg

            all_angles.append(float(angle_deg))

            # 计算尾首方向的角度
            direction_vector = start_point - end_point
            direction_angle_rad = np.arccos(np.dot(direction_vector, horizontal_vector) / (np.linalg.norm(direction_vector) * np.linalg.norm(horizontal_vector)))
            direction_angle_deg = np.degrees(direction_angle_rad)

            if direction_vector[1] < 0:
                direction_angle_deg = 360 - direction_angle_deg

            all_direction_angles.append(float(direction_angle_deg))

        return (all_angles, all_direction_angles)





NODE_CLASS_MAPPINGS = {
    "MaskToPoint": MaskToPoint,
    "KeypointSelector": KeypointSelector,
    "CenterPointCalculator": CenterPointCalculator,
    "PointMerger": PointMerger,
    "PointExtractor": PointExtractor,
    "PointMover": PointMover,
    "PointReversor": PointReversor,
    "DistanceCalculator": DistanceCalculator,
    "PointPreview": PointPreview,
    "PointToCoordinates": PointToCoordinates,
    "CoordinatesToPoint": CoordinatesToPoint,
    "PointThresholdFilter": PointThresholdFilter,
    "GridPointGenerator": GridPointGenerator,
    "AngleCalculator": AngleCalculator,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MaskToPoint": "Mask to Point",
    "KeypointSelector": "Keypoint Selector",
    "CenterPointCalculator": "Center Point Calculator",
    "PointMerger": "Point Merger",
    "PointExtractor": "Point Extractor",
    "PointMover": "Point Mover",
    "PointReversor": "Point Reversor",
    "DistanceCalculator": "Distance Calculator",
    "PointPreview": "Point Preview",
    "PointToCoordinates": "Point to Coordinates",
    "CoordinatesToPoint": "Coordinates to Point",
    "PointThresholdFilter": "Point Threshold Filter",
    "GridPointGenerator": "Grid Point Generator",
    "AngleCalculator": "Angle Calculator",
}


